import argparse
import sys

import yaml

args = None
# config_fn = './experiments/config_conll03_wm.yaml'
# config_fn = './experiments/config_splitcifar100_supsup.yaml'
# config_fn = './experiments/config_splitcifar100_dst.yaml'
config_fn = './experiments/config_splitcifar100_wm.yaml'


def parse_arguments():
    # Training settings
    parser = argparse.ArgumentParser(description="MTL_mw")
    parser.add_argument("--debug", type=bool, default=True)
    parser.add_argument("--seed", type=int, default=1)

    # io path & args
    parser.add_argument("--task_lst_name",
                        type=str,
                        default='./data/conll03.pkl',
                        help=".pkl which contians task list")
    parser.add_argument("--log-dir",
                        type=str,
                        default='./runs/',
                        help="Location to logs/checkpoints")
    parser.add_argument('--exp_name',
                        type=str,
                        default='exp1',
                        help='the experiment name')
    parser.add_argument("--save",
                        action="store_true",
                        default=False,
                        help="save checkpoints")
    parser.add_argument("--save_model_path",
                        type = str,
                        default = None,
                        help = "save checkpoints")

    # model
    parser.add_argument("--model",
                        type=str,
                        default='cnn-lstm',
                        help="Type of model.")
    parser.add_argument("--num_tasks", type=int, default=3, help="# of tasks")
    parser.add_argument('--hidden_size',
                        type=int,
                        default=200,
                        help='hidden size of encoder')
    parser.add_argument('--n_layer',
                        type=int,
                        default=1,
                        help='num of layers of encoder')
    parser.add_argument('--crf',
                        action='store_true',
                        default=False,
                        help='whether use crf')
    parser.add_argument("--dropout",
                        dest='dropout',
                        type=float,
                        default=0.5,
                        help='the dropout probability')

    # training
    parser.add_argument("--trainer",
                        type=str,
                        default='trainer_task_first',
                        help="Type of trainer.")
    parser.add_argument("--alpha", type=float, default=1e-6)
    parser.add_argument("--epochs",
                        type=int,
                        default=100,
                        metavar="N",
                        help="number of epochs to train (default: 100)")
    parser.add_argument("--train_weight_lr",
                        type=float,
                        default=0.01,
                        help="learning rate (default: 0.01)")
    parser.add_argument("--train_score_lr",
                        type=float,
                        default=0.1,
                        help="learning rate (default: 0.1)")
    parser.add_argument("--momentum",
                        type=float,
                        default=0.9,
                        metavar="M",
                        help="Momentum (default: 0.9)")
    parser.add_argument("--wd",
                        type=float,
                        default=0.0001,
                        metavar="M",
                        help="Weight decay (default: 0.0001)")
    parser.add_argument("--optimizer",
                        type=str,
                        default="sgd",
                        help="Which optimizer to use")
    parser.add_argument("--batch-size",
                        type=int,
                        default=128,
                        metavar="N",
                        help="input batch size for training (default: 64)")
    parser.add_argument("--test-batch-size",
                        type=int,
                        default=128,
                        metavar="N",
                        help="input batch size for testing (default: 128)")

    # visualization
    parser.add_argument('--accumulation_steps', type=int, default=1)
    parser.add_argument('--print_every',
                        type=int,
                        default=1,
                        help='the logging intervals')

    # others
    parser.add_argument("--config",
                        type=str,
                        default=None,
                        help="Config file to use, YAML format")
    parser.add_argument("--workers",
                        type=int,
                        default=4,
                        help="how many cpu workers")
    parser.add_argument("--multigpu",
                        default=None,
                        type=lambda x: [int(a) for a in x.split(",")],
                        help="Which GPUs to use for multigpu training")

    parser.add_argument(
        "--sparsity",
        type=float,
        default=0.5,
        help="how sparse is each layer, when using MultitaskMaskConv")

    parser.add_argument(
        "--pruning_method",
        default=None,
        type=str,
        help="choose pruning method from 'topK' 'threshold' 'sigmoied_threshold' 'l0'"
    )
    parser.add_argument(
        "--w_epochs",
        default=None,
        type=int,
        help="weight 训练多少次'"
    )
    parser.add_argument(
        "--lr_epochs",
        default=None,
        type=int,
        help="开始第二次schdular的epoch数，即保持lr的拖尾"
    )
    parser.add_argument(
        "--num_params",
        default=None,
        type=list,
        help="模型包括多少有效的参数"
    )
    parser.add_argument(
        "--lr1_min",
        type=float,
        default=None,
        help="第一次Lr下降策略下最小的lr值"
    )
    parser.add_argument(
        "--init_method",
        default='kaiming',
        type=str,
        help="choose init method from 'kaiming' 'lr_constant' "
    )
    parser.add_argument(
        "--alpha_epoch",
        default=170,
        type=int,
        help="拖尾的alpha是从哪个epoch开始 "
    )
    parser.add_argument(
        "--alpha_epoch_1",
        default=None,
        type=int,
        help="多少个epoch为一组波峰"
    )
    parser.add_argument(
        "--alpha_epoch_2",
        default=None,
        type=int,
        help="每个波峰中lr=lr的前多少个"
    )
    parser.add_argument(
        "--intercept",
        default=0.02,
        type=float,
        help="loss v4中，需要调的截距项"
    )


    # sparing sharing
    parser.add_argument("--final_rate", default=0.1, type=float, help="spare sharing mask 迭代的最小值")
    parser.add_argument("--prune_iter", default=20, type=int, help="迭代到最小Mask的次数")
    parser.add_argument("--warmup_save", default=False, type=bool, help="是否需要储存warmup")
    parser.add_argument("--update_mask", default=False, type=bool, help="在ss多任务训练时是否更新Mask")
    parser.add_argument("--pruning_iter", default=5, type=int, help="在ss多任务训练时,选择多少次迭代的mask")

    parser.add_argument("--task_id", default=-1, type=int, help="任务id，应用于task frist")
    parser.add_argument("--tasks", default=[0,1,2], type=list, help="任务id，哪些任务id要被训练")
    # 多任务
    parser.add_argument("--scheduler_weight", default=False, type=bool, help="在多任务中是否需要weight的scheduler")
    parser.add_argument("--scheduler_score", default=True, type=bool, help="在多任务中是否需要mask的scheduler")
    parser.add_argument("--threshold_pruning", default=True, type=bool, help="是否使用threshold剪枝的方法,否则为topk")

    args = parser.parse_args()

    # Allow for use from notebook without config file
    if args.config is not None:
        get_config(args)

    return args


def get_config(args):
    # get commands from command line
    override_args = argv_to_vars(sys.argv)

    # load yaml file
    yaml_txt = open(args.config).read()

    # override args
    loaded_yaml = yaml.load(yaml_txt, Loader=yaml.FullLoader)
    for v in override_args:
        loaded_yaml[v] = getattr(args, v)

    # print(f"=> Reading YAML config from {args.config}")
    args.__dict__.update(loaded_yaml)


def trim_preceding_hyphens(st):
    i = 0
    while st[i] == "-":
        i += 1

    return st[i:]


def arg_to_varname(st: str):
    st = trim_preceding_hyphens(st)
    st = st.replace("-", "_")

    return st.split("=")[0]


def argv_to_vars(argv):
    var_names = []
    for arg in argv:
        if arg.startswith("-") and arg_to_varname(arg) != "config":
            var_names.append(arg_to_varname(arg))

    return var_names


def get_args():
    global args
    if args is None:
        args = parse_arguments()


def from_config_file(fn=config_fn):
    # load yaml file
    yaml_txt = open(fn).read()

    loaded_yaml = yaml.load(yaml_txt, Loader=yaml.FullLoader)

    # print(f"=> Reading YAML config from {fn}")

    global args
    parser = argparse.ArgumentParser(description="MTL_mw")
    args = parser.parse_args()
    for k in loaded_yaml:
        # setattr(args, k, loaded_yaml[k])
        args.__dict__[k] = loaded_yaml[k]


# from_config_file()
get_args()
